# Copyright (c) 2010
# The Regents of the University of Michigan
# All Rights Reserved

# Permission is granted to use, copy, create derivative works, and
# redistribute this software and such derivative works for any purpose,
# so long as the name of the University of Michigan is not used in
# any advertising or publicity pertaining to the use or distribution
# of this software without specific, written prior authorization. If
# the above copyright notice or any other identification of the
# University of Michigan is included in any copy of any portion of
# this software, then the disclaimer below must also be included.

# This software is provided as is, without representation or warranty
# of any kind either express or implied, including without limitation
# the implied warranties of merchantability, fitness for a particular
# purpose, or noninfringement.  The Regents of the University of
# Michigan shall not be liable for any damages, including special,
# indirect, incidental, or consequential damages, with respect to any
# claim arising out of or in connection with the use of the software,
# even if it has been or is hereafter advised of the possibility of
# such damages.

'''Python impl of serialization library for Vuze.

Mostly dead due to Cython re-implementation, but acts as a fallback when
something hasn\'t been implemented in Cython yet.

Fair warning: make sure that the serialization functions you want to
use have actually been implemented. I only bothered to implement
things I actually needed for crawling the DHT. In particular, at least
one reader will be disappointed to note that this code completely
ignores IPv6.
'''

import collections
import logging
import socket
import struct
import time

MIN_SUPPORTED_PROTOCOL_VERSION = 22
SUPPORTED_PROTOCOL_VERSION = 22


# DHT packet type codes.
ACT_REQUEST_PING = 1024
ACT_REPLY_PING = 1025
ACT_REQUEST_STORE = 1026
ACT_REPLY_STORE = 1027
ACT_REQUEST_FIND_NODE = 1028
ACT_REPLY_FIND_NODE = 1029
ACT_REQUEST_FIND_VALUE = 1030
ACT_REPLY_FIND_VALUE = 1031
ACT_REPLY_ERROR = 1032
ACT_REPLY_STATS = 1033
ACT_REQUEST_STATS = 1034
ACT_DATA = 1035


# For DATA packets.
(PT_READ_REQUEST, PT_READ_REPLY, PT_WRITE_REQUEST, PT_WRITE_REPLY) = range(4)


##### Lots of structures #####

HARMLESS_VIVALDI = ("\x01\x01\x10" +  # 1 position, type 1, sz16
  "\x00\x00\x00\x00" +  # x-coord
  "\x00\x00\x00\x00" +  # y-coord
  "\x00\x00\x00\x00" +  # height
  "\x7f\x80\x00\x00")    # error (+infty)


# Remember to byteswap contact addresses again before feeding to inet_ntoa.
RECV_REQUEST_HEADER = struct.Struct('!QIIBBIBxIHIQ')
SEND_REQUEST_HEADER = struct.Struct('!QIIBBIBBIHIQ')


REPLY_HEADER = struct.Struct('!IIQBBII')


CONTACT_ADDRESS = struct.Struct('!xIH')


RECV_CONTACT = struct.Struct('!xxxIH')
SEND_CONTACT = struct.Struct('!BBBIH')

REQUEST_STORE_BODY_START = struct.Struct('!IB')


# ID length ignored on recv, but we have to send it.
RECV_REQUEST_FIND_NODE_BODY = struct.Struct('!x20s8x')
SEND_REQUEST_FIND_NODE_BODY = struct.Struct('!B20sII')


RECV_REPLY_FIND_NODE_BODY_START = struct.Struct('!IIIB')
SEND_REPLY_FIND_NODE_BODY_START = struct.Struct('!III')


# ID length ignored on recv, but we have to send it.
RECV_REQUEST_FIND_VALUE_BODY = struct.Struct('!x20sBB')
SEND_REQUEST_FIND_VALUE_BODY = struct.Struct('!B20sBB')


REPLY_FIND_VALUE_BODY_START = struct.Struct('!BB')



RECV_VIVALDI_HEADER = struct.Struct('!xB')


RECV_VALUE_START = struct.Struct('!I8xH')  # Up to the value's bytes.
SEND_VALUE_START = struct.Struct('!IQH')


DATA_START = struct.Struct('!BB') # type, length
DATA_TRAILER = struct.Struct('!IIIH')


REPLY_STATS_START = struct.Struct('!17Q')
REPLY_STATS_TRAILER = struct.Struct('!QI2Q5Q')
KEY = struct.Struct('!x20s')
SEND_KEY = struct.Struct('!B20s')


# Diversification types.
DT_NONE, DT_FREQ, DT_SIZE = range(1, 4)

##### Deserialized packets #####

RequestHeader = collections.namedtuple('RequestHeader', ['connection_id',
                                                         'action',
                                                         'transaction_id',
                                                         'protocol_version',
                                                         'vendor_id',
                                                         'network_id',
                                                         'local_protocol_version',
                                                         'node_address',
                                                         'instance_id',
                                                         'timestamp'])


ReplyHeader = collections.namedtuple('ReplyHeader', ['action',
                                                     'transaction_id',
                                                     'connection_id',
                                                     'protocol_version',
                                                     'vendor_id',
                                                     'network_id',
                                                     'instance_id'])


RequestPing = collections.namedtuple('RequestPing', ['header'])


ReplyPing = collections.namedtuple('ReplyPing', ['header'])


RequestStore = collections.namedtuple('RequestStore', ['header',
                                                       'spoof_id',
                                                       'keys',
                                                       'value_groups'])


ReplyStore = collections.namedtuple('ReplyStore', ['header',
                                                   'diversifications'])


RequestFindNode = collections.namedtuple('RequestFindNode', ['header', 'id'])


ReplyFindNode = collections.namedtuple('ReplyFindNode', ['header',
                                                         'spoof_id',
                                                         'node_type',
                                                         'dht_size',
                                                         'contacts'])


RequestFindValue = collections.namedtuple('RequestFindValue', ['header',
                                                               'id',
                                                               'flags',
                                                               'max_values'])


ReplyFindValue = collections.namedtuple('ReplyFindValue', ['header',
                                                           'has_continuation',
                                                           'contacts',
                                                           'diversification_type',
                                                           'values'])


# We'll never want to request the NP_VER_2 stuff, so don't expose stats type.
RequestStats = collections.namedtuple('RequestStats', ['header'])

ReplyStats = collections.namedtuple('ReplyStats', ['header',
                                                   'db_values_stored',
                                                   'router_nodes',
                                                   'router_leaves',
                                                   'router_contacts',
                                                   'total_bytes_received',
                                                   'total_bytes_sent',
                                                   'total_packets_received',
                                                   'total_packets_sent',
                                                   'total_pings_received',
                                                   'total_find_nodes_received',
                                                   'total_find_values_received',
                                                   'total_stores_received',
                                                   'average_bytes_received',
                                                   'average_bytes_sent',
                                                   'average_packets_received',
                                                   'average_packets_sent',
                                                   'incoming_requests',
                                                   'az_version',
                                                   'router_uptime',
                                                   'router_count',
                                                   'db_keys_blocked',
                                                   'total_key_blocks_received',
                                                   'db_key_count',
                                                   'db_value_count',
                                                   'db_store_size',
                                                   'db_freq_divs',
                                                   'db_size_divs'])


# Looks like transfer_key is some ridiculous implementation-specific
# mechanism (used to find the callback that should handle this
# request); we'll need to sniff packets, find one of the kind we want
# to send, and hardcode the outgoing transfer_key. (UPDATE: easier
# than that, see below)

# Reference org/gudy/azureus2/pluginsimpl/local/ddb/DDBaseImpl.java
# specifically the line where it passes
# "DDBaseHelpers.getKey(type.getClass()).getHash()," as the type. The
# class in question for our purposes is DDBaseTTTorrent, and the hash
# function is actually just the SHA1 of the class name!
Data = collections.namedtuple('Data', ['header',
                                       'packet_type',
                                       'transfer_key',
                                       'key',
                                       'start',
                                       'length',
                                       'total_length',
                                       'data'])


TransportValue = collections.namedtuple('TransportValue', ['version',
                                                           'bytes',
                                                           'originator',
                                                           'flags'])



##### Actual implemention #####


class SerializationError(Exception):
  pass


def decode_fixup_ip_addr(addr):
  '''Compensate for extra, unnecessary byteswapping on IP address decoding.'''
  addr = struct.pack('!I', addr)
  return socket.inet_ntoa(addr)


def encode_fixup_ip_addr(addr):
  '''Compensate for extra, unnecessary byteswapping on IP address decoding.'''
  return struct.unpack('!I', socket.inet_aton(addr))[0]


def decode_request_header(packet):
  res = RECV_REQUEST_HEADER.unpack(packet[:RECV_REQUEST_HEADER.size])
  node_address = (decode_fixup_ip_addr(res[-4]), res[-3])
  res = res[:-4] + (node_address,) + res[-2:]
  return (RequestHeader(*res), packet[RECV_REQUEST_HEADER.size:])


def decode_reply_header(packet):
  res = REPLY_HEADER.unpack(packet[:REPLY_HEADER.size])
  return (ReplyHeader(*res), packet[REPLY_HEADER.size:])


def encode_request_header(req_header):
  addr = req_header[-3]
  rep = (req_header[:-3] + (4, encode_fixup_ip_addr(addr[0]), addr[1]) +
         req_header[-2:])
  return SEND_REQUEST_HEADER.pack(*rep)


def encode_reply_header(header):
  return REPLY_HEADER.pack(*header)


def decode_contact(packet):
  if not packet:
    raise SerializationError('Empty packet passed to decode_contact')

  if ord(packet[2]) == 16:
    logging.info('Ignored IPv6 contact.')
    return (None, packet[21:])
  ip, port = RECV_CONTACT.unpack(packet[:RECV_CONTACT.size])
  ip = decode_fixup_ip_addr(ip)
  packet = packet[RECV_CONTACT.size:]
  return ((ip, port), packet)


def encode_contact(address):
  return SEND_CONTACT.pack(1, SUPPORTED_PROTOCOL_VERSION, 4,
                           encode_fixup_ip_addr(address[0]),
                           address[1])

def skip_vivaldi(packet, count):
  for x in xrange(count):
    size = RECV_VIVALDI_HEADER.unpack(packet[:RECV_VIVALDI_HEADER.size])[0]
    packet = packet[(RECV_VIVALDI_HEADER.size + size):]
  return packet


def encode_ping_request(request):
  return encode_request_header(request.header)


def encode_ping_reply(reply):
  header = encode_reply_header(reply.header)
  return header + HARMLESS_VIVALDI


def decode_value_group(packet):
  value_group = []
  values_count = struct.unpack('!H', packet[:2])[0]
  packet = packet[2:]
  for y in xrange(values_count):
    (version, value_bytes_count) = RECV_VALUE_START.unpack(
      packet[:RECV_VALUE_START.size])
    packet = packet[RECV_VALUE_START.size:]
    assert value_bytes_count >= 0
    if value_bytes_count:
      bytes = packet[:value_bytes_count]
    else:
      bytes = str()
    originator, packet = decode_contact(packet[value_bytes_count:])
    flags = ord(packet[0])
    value_group.append(TransportValue(version, bytes, originator, flags))
    packet = packet[1:]  # Skip flags byte!

  # TODO: there is extra crap at the end for newer protocol versions. If we
  # go to them, we'll need to support it.

  return (value_group, packet)


def encode_value_group(vg):
  values_count = struct.pack('!H', len(vg))

  ev = encode_value
  encoded_packet = [ev(value) for value in vg]
  return ''.join((values_count, ''.join(encoded_packet)))


def encode_value(value):
  encoded_start = SEND_VALUE_START.pack(value.version,
                                        int(time.time() * 1000),
                                        len(value.bytes))
  encoded_chunk = [encoded_start, value.bytes,
                   encode_contact(value.originator), chr(value.flags)]
  # TODO: extra crap for newer proto versions, see above TODO.
  return ''.join(encoded_chunk)
                                        

def encode_store_request(request):
  header = encode_request_header(request.header)
  encoded_packet = [header, REQUEST_STORE_BODY_START.pack(request.spoof_id,
                                                  len(request.keys))]

  for key in request.keys:
    encoded_packet.append(SEND_KEY.pack(20, key))

  encoded_packet.append(chr(len(request.value_groups)))

  for vg in request.value_groups:
    encoded_packet.append(encode_value_group(vg))
  return ''.join(encoded_packet)


def decode_store_request(packet, header):
  (spoof_id, keys_count) = REQUEST_STORE_BODY_START.unpack(
      packet[:REQUEST_STORE_BODY_START.size])
  packet = packet[REQUEST_STORE_BODY_START.size:]
  keys = []
  for x in xrange(keys_count):
    try:
      keys.append(KEY.unpack(packet[:KEY.size])[0])
      packet = packet[KEY.size:]
    except struct.error:
      print x, keys_count, repr(packet)
      raise

  value_groups_count = ord(packet[0])
  packet = packet[1:]

  value_groups = []
  for x in xrange(value_groups_count):
    value_group, packet = decode_value_group(packet)
    value_groups.append(value_group)

  return RequestStore(header, spoof_id, keys, value_groups)
  

def encode_store_reply(reply):
  header = encode_reply_header(reply.header)
  return header + '\x00'


def decode_store_reply(packet, header):
  # TODO: does anyone care about diversifications?
  return ReplyStore(header, [])


def encode_find_node_request(request):
  header = encode_request_header(request.header)
  return (header + SEND_REQUEST_FIND_NODE_BODY.pack(20, request.id, 0xFFFFFFFF, 0))
  

def decode_find_node_request(packet, header):
  if packet[0] == '\x00':  # these messages are for anti-spoof.
    return RequestFindNode(header, None)
  id = RECV_REQUEST_FIND_NODE_BODY.unpack(
    packet[:RECV_REQUEST_FIND_NODE_BODY.size])[0]
  return RequestFindNode(header, id)


def encode_find_node_reply(reply):
  header = encode_reply_header(reply.header)
  packet = [header]
  packet.append(SEND_REPLY_FIND_NODE_BODY_START.pack(reply.spoof_id,
                                                     reply.node_type,
                                                     reply.dht_size))
  packet.append(HARMLESS_VIVALDI)
  packet.append(struct.pack('!H', len(reply.contacts)))

  for ip, port in reply.contacts:
    packet.append('\x01%s\x04' % chr(SUPPORTED_PROTOCOL_VERSION))
    packet.append(socket.inet_aton(ip))
    packet.append(struct.pack('!H', port))

  return ''.join(packet)


def decode_find_node_reply(packet, header):
  (spoof_id,
   node_type,
   dht_size,
   vivaldi_count) = RECV_REPLY_FIND_NODE_BODY_START.unpack(
      packet[:RECV_REPLY_FIND_NODE_BODY_START.size])
  packet = skip_vivaldi(packet[RECV_REPLY_FIND_NODE_BODY_START.size:],
                        vivaldi_count)
  
  contacts_count = struct.unpack('!H', packet[:2])[0]
  packet = packet[2:]
  contacts = []
  for x in xrange(contacts_count):
    contact, packet = decode_contact(packet)
    if contact is not None:
      contacts.append(contact)

  return ReplyFindNode(header, spoof_id, node_type, dht_size, contacts)


def decode_find_value_request(packet, header):
  (id, flags, max_values) = REPLY_FIND_VALUE_BODY.unpack(packet)
  return RequestFindValue(header, id, flags, max_values)


def encode_find_value_request(request):
  header = encode_request_header(request.header)
  return header + SEND_REQUEST_FIND_VALUE_BODY.pack(20, request.id,
                                                   request.flags,
                                                   request.max_values)


def decode_find_value_reply(packet, header):
  unpacker = REPLY_FIND_VALUE_BODY_START
  (has_continuation,
   has_values) = unpacker.unpack(packet[:unpacker.size])
  packet = packet[unpacker.size:]

  contacts = []
  values = []
  div_type = DT_NONE
  if not has_values:
    contacts_count = struct.unpack('!H', packet[:2])[0]
    packet = packet[2:]
    for _ in xrange(contacts_count):
      contact, packet = decode_contact(packet)
      if contact is not None:
        contacts.append(contact)
    # There are also network coordinates, but I'm too lazy to implement that.
  else:
    div_type = ord(packet[0])
    packet = packet[1:]
    values, packet = decode_value_group(packet)
    assert not packet, 'extra crap on the end of find_value reply!'

  return ReplyFindValue(header, has_continuation, contacts,
                        div_type, values)


def encode_find_value_reply(reply):
  # TODO: too lazy right now.
  raise NotImplementedError()


def encode_data(packet):
  header = encode_request_header(packet.header)
  return (header + chr(packet.packet_type) + chr(len(packet.transfer_key))
          + packet.transfer_key + chr(len(packet.key)) + packet.key
          + DATA_TRAILER.pack(packet.start, packet.length,
                              packet.total_length, len(packet.data))
          + packet.data)


def decode_data(packet, header):
  (packet_type,
   transfer_key_length) = DATA_START.unpack(packet[:DATA_START.size])

  packet = packet[DATA_START.size:]
  transfer_key = packet[:transfer_key_length]
  packet = packet[transfer_key_length:]

  key_length = ord(packet[0])
  key = packet[1:key_length + 1]
  packet = packet[key_length + 1:]
  (start,
   length,
   total_length,
   data_length) = DATA_TRAILER.unpack(packet[:DATA_TRAILER.size])
  data = packet[DATA_TRAILER.size:DATA_TRAILER.size + data_length]
  return Data(header, packet_type, transfer_key, key, start, length,
              total_length, data)


def encode_stats_request(packet):
  header = encode_request_header(packet.header)
  return header + '\x00\x00\x00\x01'


def decode_stats_reply(packet, header):
  packet = packet[4:] # Don't care about type.
  stats_head = REPLY_STATS_START.unpack(packet[:REPLY_STATS_START.size])
  packet = packet[REPLY_STATS_START.size:]
  version_length = ord(packet[0])
  version = packet[1:version_length + 1]
  packet = packet[version_length + 1:]
  stats_trailer = REPLY_STATS_TRAILER.unpack(packet[:REPLY_STATS_TRAILER.size])
  return ReplyStats(stats_head + (version,) + stats_trailer)
  

# TODO: calling program may not want to do full decodes for all of these.
#       Provide a way for the caller to indicate that things need not be
#       decoded.
def decode_packet(packet, handler, *args):
  '''Deserialize packet, determine its type, and callback to handler.

  This is the intended entry point for this module.
  
  packet is expected to be a sequence of raw bytes from the network.
  handler is expected to have the same interface as dhtutil.BaseClient
  (see code).
  args is passed right along to each of the handler functions, in case
  the client needs extra arguments.
  '''
  header = None
  if not packet:
    raise SerializationError('Empty packet received!')
  try:
    if ord(packet[0]) & 0x80:
      header, packet = decode_request_header(packet)

      if header.protocol_version < 14:
        raise SerializationError('Protocol version too old!')

      if header.action == ACT_REQUEST_PING:
        handler.handle_ping_request(header, *args)
      elif header.action == ACT_REQUEST_STORE:
        request = decode_store_request(packet, header)
        handler.handle_store_request(request, *args)
      elif header.action == ACT_REQUEST_FIND_NODE:
        request = decode_find_node_request(packet, header)
        handler.handle_find_node_request(request, *args)
      elif header.action == ACT_DATA:
        request = decode_data(packet, header)
        handler.handle_data(request, *args)
    else:
      header, packet = decode_reply_header(packet)

      if header.protocol_version < 14:
        raise SerializationError('Protocol version too old!')

      if header.action == ACT_REPLY_FIND_NODE:
        reply = decode_find_node_reply(packet, header)
        handler.handle_find_node_reply(reply, *args)
      elif header.action == ACT_REPLY_PING:
        handler.handle_ping_reply(header, *args)
      elif header.action == ACT_REPLY_STATS:
        reply = decode_stats_reply(packet, header)
        handler.handle_stats_reply(reply, *args)
      elif header.action == ACT_REPLY_FIND_VALUE:
        reply = decode_find_value_reply(packet, header)
        handler.handle_find_value_reply(reply, *args)
      elif header.action == ACT_REPLY_STORE:
        reply = decode_store_reply(packet, header)
        handler.handle_store_reply(reply, *args)

  except struct.error as e:
    if header is not None:
      logging.error('Action: %d Protocol version: %d' % (header.action, header.protocol_version))
    raise SerializationError(e)
  except SerializationError as e:
    if header is not None:
      logging.debug('Action: %d' % header.action)
    raise
